import akshare as ak
import psycopg2
from psycopg2 import sql
import datetime

# 数据库连接配置
DB_CONFIG = {
    'host': 'localhost',
    'port': '5432',
    'database': 'stock_db',
    'user': 'leizhen',
    'password': '751982leizhen'
}

def get_db_connection():
    """获取数据库连接"""
    return psycopg2.connect(**DB_CONFIG)

def create_table_if_not_exists(conn):
    """创建股东分析表"""
    with conn.cursor() as cursor:
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS stock_gdfx_holding_analyse_em (
                stock_code VARCHAR(10) NOT NULL,
                date DATE NOT NULL,
                shareholder_name VARCHAR(100),
                shareholder_type VARCHAR(50),
                share_number NUMERIC(20,4),
                share_ratio NUMERIC(20,4),
                change_type VARCHAR(20),
                change_number NUMERIC(20,4),
                change_ratio NUMERIC(20,4),
                PRIMARY KEY (stock_code, date, shareholder_name)
            );
        """)
        conn.commit()

def get_quarter_end_dates(start_year=2015):
    """生成季度末日期列表"""
    end_year = datetime.datetime.now().year
    quarter_ends = []
    
    for year in range(start_year, end_year + 1):
        for month in [3, 6, 9, 12]:
            # 获取季度最后一天
            if month == 12 or month == 3:
                quarter_end = datetime.date(year, month, 31)
            else:  
                quarter_end = datetime.date(year, month, 30)
            # 只添加过去的日期
            if quarter_end < datetime.date.today():
                quarter_ends.append(quarter_end.strftime("%Y%m%d"))
    
    return quarter_ends

def save_to_database(conn, df, date):
    """保存数据到数据库"""
    if df is None or df.empty:
        return
    
    # 打印列名用于调试
    print("原始数据列名:", df.columns.tolist())
    
    # 检查并重命名列名
    column_mapping = {
        '股票代码': '股票代码',
        '股东名称': '股东名称',
        '期末持股-数量': 'share_number',
        '期末持股-数量变化': 'change_number',
        '期末持股-数量变化比例': 'change_ratio',
        '期末持股-持股变动': 'change_type',
        '占总流通股本持股比例': 'share_ratio'
    }
    
    # 重命名列名
    for old_name, new_name in column_mapping.items():
        if old_name in df.columns:
            df = df.rename(columns={old_name: new_name})
    
    # 检查必要列是否存在
    required_columns = ['股票代码', '股东名称']
    missing_cols = [col for col in required_columns if col not in df.columns]
    if missing_cols:
        print(f"缺少必要列: {missing_cols}")
        return
    
    with conn.cursor() as cursor:
        for _, row in df.iterrows():
            try:
                # 准备插入数据
                insert_sql = sql.SQL("""
                    INSERT INTO stock_gdfx_holding_analyse_em (
                        stock_code, date, shareholder_name,
                        share_number, share_ratio, change_type, change_number, change_ratio
                    ) VALUES (
                        %s, %s, %s, %s, %s, %s, %s, %s
                    )
                    ON CONFLICT (stock_code, date, shareholder_name) DO UPDATE SET
                        share_number = EXCLUDED.share_number,
                        share_ratio = EXCLUDED.share_ratio,
                        change_type = EXCLUDED.change_type,
                        change_number = EXCLUDED.change_number,
                        change_ratio = EXCLUDED.change_ratio
                """)
                
                # 打印当前行数据用于调试
                # print("当前行数据:", row.to_dict())
                
                cursor.execute(insert_sql, (
                    row['股票代码'],
                    datetime.datetime.strptime(date, "%Y%m%d").date(),
                    row['股东名称'],
                    row.get('share_number', 0),
                    row.get('share_ratio', 0),
                    row.get('change_type', ''),
                    row.get('change_number', 0),
                    row.get('change_ratio', 0)
                ))
                conn.commit()  # 每次插入后立即提交
            except Exception as e:
                print(f"保存数据失败: {e}")
                conn.rollback()

def main():
    conn = get_db_connection()
    create_table_if_not_exists(conn)
    
    quarter_dates = get_quarter_end_dates()
    
    for date in quarter_dates:
        try:
            print(f"正在获取 {date} 的股东持股数据...")
            df = ak.stock_gdfx_holding_analyse_em(date=date)
            save_to_database(conn, df, date)
            print(f"{date} 数据保存成功")
        except Exception as e:
            print(f"获取 {date} 数据失败: {e}")
    
    conn.close()

if __name__ == "__main__":
    main()
